﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Runtime.InteropServices;
using SDFL.Native;
using System.Diagnostics;
using SDFL.Core;
using Microsoft.Win32.SafeHandles;

namespace SDFL.Helper
{
    public class FileDetails
    {
        public string Name { get; set; }
    }

    public class HandleHelper
    {

        static int m_nLastProcId = 0;
        
        public static List<NativeMethods.SYSTEM_HANDLE_INFORMATION> GetHandles(System.Diagnostics.Process process)
        {
            uint nStatus;
            int nHandleInfoSize = 0x10000;
            IntPtr ipHandlePointer = Marshal.AllocHGlobal(nHandleInfoSize);
            int nLength = 0;
            IntPtr ipHandle = IntPtr.Zero;

            while ((nStatus = NativeMethods.NtQuerySystemInformation(NativeMethods.CNST_SYSTEM_HANDLE_INFORMATION, ipHandlePointer, nHandleInfoSize, ref nLength)) == NativeMethods.STATUS_INFO_LENGTH_MISMATCH)
            {
                nHandleInfoSize = nLength;
                Marshal.FreeHGlobal(ipHandlePointer);
                ipHandlePointer = Marshal.AllocHGlobal(nLength);
            }

            byte[] baTemp = new byte[nLength];
            NativeMethods.CopyMemory(baTemp, ipHandlePointer, (uint)nLength);

            long lHandleCount = 0;
            if (SDFL.Core.SystemInfo.Is64Bits())
            {
                lHandleCount = Marshal.ReadInt64(ipHandlePointer);
                ipHandle = new IntPtr(ipHandlePointer.ToInt64() + 8);
            }
            else
            {
                lHandleCount = Marshal.ReadInt32(ipHandlePointer);
                ipHandle = new IntPtr(ipHandlePointer.ToInt32() + 4);
            }

            NativeMethods.SYSTEM_HANDLE_INFORMATION shHandle;
            List<NativeMethods.SYSTEM_HANDLE_INFORMATION> lstHandles = new List<NativeMethods.SYSTEM_HANDLE_INFORMATION>();

            for (long lIndex = 0; lIndex < lHandleCount; lIndex++)
            {
                shHandle = new NativeMethods.SYSTEM_HANDLE_INFORMATION();
                ConvertIpHandleShHandle(ref ipHandle, ref shHandle);
                if (shHandle.ProcessID != process.Id) continue;
                lstHandles.Add(shHandle);
            }
            return lstHandles;

        }        

        public static FileDetails GetFileDetails(IntPtr processHwnd, NativeMethods.SYSTEM_HANDLE_INFORMATION sYSTEM_HANDLE_INFORMATION)
        {
            FileDetails fd = new FileDetails();
            fd.Name = "";
            IntPtr ipHandle = IntPtr.Zero;
            NativeMethods.OBJECT_BASIC_INFORMATION objBasic = new NativeMethods.OBJECT_BASIC_INFORMATION();
            IntPtr ipBasic = IntPtr.Zero;
            NativeMethods.OBJECT_TYPE_INFORMATION objObjectType = new NativeMethods.OBJECT_TYPE_INFORMATION();
            IntPtr ipObjectType = IntPtr.Zero;
            NativeMethods.OBJECT_NAME_INFORMATION objObjectName = new NativeMethods.OBJECT_NAME_INFORMATION();
            IntPtr ipObjectName = IntPtr.Zero;
            string strObjectTypeName = "";
            string strObjectName = "";
            int nLength = 0;
            int nReturn = 0;
            IntPtr ipTemp = IntPtr.Zero;

            OpenProcessForHandle(processHwnd, sYSTEM_HANDLE_INFORMATION.ProcessID);            
            if (!NativeMethods.DuplicateHandle(processHwnd, sYSTEM_HANDLE_INFORMATION.Handle, NativeMethods.GetCurrentProcess(), out ipHandle, 0, false, NativeMethods.DUPLICATE_SAME_ACCESS)) return fd;

            ipBasic = Marshal.AllocHGlobal(Marshal.SizeOf(objBasic));
            NativeMethods.NtQueryObject(ipHandle, (int)NativeMethods.ObjectInformationClass.ObjectBasicInformation, ipBasic, Marshal.SizeOf(objBasic), ref nLength);
            objBasic = (NativeMethods.OBJECT_BASIC_INFORMATION)Marshal.PtrToStructure(ipBasic, objBasic.GetType());
            Marshal.FreeHGlobal(ipBasic);


            ipObjectType = Marshal.AllocHGlobal(objBasic.TypeInformationLength);
            nLength = objBasic.TypeInformationLength;
            while ((uint)(nReturn = NativeMethods.NtQueryObject(ipHandle, (int)NativeMethods.ObjectInformationClass.ObjectTypeInformation, ipObjectType, nLength, ref nLength)) == NativeMethods.STATUS_INFO_LENGTH_MISMATCH)
            {
                Marshal.FreeHGlobal(ipObjectType);
                ipObjectType = Marshal.AllocHGlobal(nLength);
            }

            objObjectType = (NativeMethods.OBJECT_TYPE_INFORMATION)Marshal.PtrToStructure(ipObjectType, objObjectType.GetType());
            if (SystemInfo.Is64Bits())
            {
                ipTemp = new IntPtr(Convert.ToInt64(objObjectType.Name.Buffer.ToString(), 10) >> 32);
            }
            else
            {
                ipTemp = objObjectType.Name.Buffer;
            }

            strObjectTypeName = Marshal.PtrToStringUni(ipTemp, objObjectType.Name.Length >> 1);
            Marshal.FreeHGlobal(ipObjectType);
            if (strObjectTypeName != "File") return fd;

            nLength = objBasic.NameInformationLength;
            ipObjectName = Marshal.AllocHGlobal(nLength);
            while ((uint)(nReturn = NativeMethods.NtQueryObject(ipHandle, (int)NativeMethods.ObjectInformationClass.ObjectNameInformation, ipObjectName, nLength, ref nLength)) == NativeMethods.STATUS_INFO_LENGTH_MISMATCH)
            {
                Marshal.FreeHGlobal(ipObjectName);
                ipObjectName = Marshal.AllocHGlobal(nLength);
            }
            objObjectName = (NativeMethods.OBJECT_NAME_INFORMATION)Marshal.PtrToStructure(ipObjectName, objObjectName.GetType());

            if (SystemInfo.Is64Bits())
            {
                ipTemp = new IntPtr(Convert.ToInt64(objObjectName.Name.Buffer.ToString(), 10) >> 32);
            }
            else
            {
                ipTemp = objObjectName.Name.Buffer;
            }

            byte[] baTemp = new byte[nLength];
            NativeMethods.CopyMemory(baTemp, ipTemp, (uint)nLength);

            if (SystemInfo.Is64Bits())
            {
                strObjectName = Marshal.PtrToStringUni(new IntPtr(ipTemp.ToInt64()));
            }
            else
            {
                strObjectName = Marshal.PtrToStringUni(new IntPtr(ipTemp.ToInt32()));
            }

            Marshal.FreeHGlobal(ipObjectName);
            NativeMethods.CloseHandle(ipHandle);

            fd.Name = GetRegularFileNameFromDevice(strObjectName);
            return fd;
        }

        public static string GetRegularFileNameFromDevice(string strRawName)
        {
            string strFileName = strRawName;
            foreach (string strDrivePath in Environment.GetLogicalDrives())
            {
                StringBuilder sbTargetPath = new StringBuilder(NativeMethods.MAX_PATH);
                if (NativeMethods.QueryDosDevice(strDrivePath.Substring(0, 2), sbTargetPath, NativeMethods.MAX_PATH) == 0)
                {
                    return strRawName;
                }
                string strTargetPath = sbTargetPath.ToString();
                if (strFileName.StartsWith(strTargetPath))
                {
                    strFileName = strFileName.Replace(strTargetPath, strDrivePath.Substring(0, 2));
                    break;
                }
            }
            return strFileName;
        }

        private static void OpenProcessForHandle(IntPtr processHwnd, int p)
        {
            if (p != m_nLastProcId)
            {
                NativeMethods.CloseHandle(processHwnd);
                processHwnd = NativeMethods.OpenProcess(NativeMethods.ProcessAccessFlags.DupHandle, false, p);
                m_nLastProcId = p;
            }
        }

        public static void ConvertIpHandleShHandle(ref IntPtr ipHandle, ref NativeMethods.SYSTEM_HANDLE_INFORMATION shHandle)
        {
            if (SDFL.Core.SystemInfo.Is64Bits())
            {
                shHandle = (NativeMethods.SYSTEM_HANDLE_INFORMATION)Marshal.PtrToStructure(ipHandle, shHandle.GetType());
                ipHandle = new IntPtr(ipHandle.ToInt64() + Marshal.SizeOf(shHandle) + 8);
            }
            else
            {
                ipHandle = new IntPtr(ipHandle.ToInt64() + Marshal.SizeOf(shHandle));
                shHandle = (NativeMethods.SYSTEM_HANDLE_INFORMATION)Marshal.PtrToStructure(ipHandle, shHandle.GetType());
            }
        }

        /// <summary>
        /// Protect file by duplicate its handle to a stable process
        /// </summary>
        /// <param name="filePath">the file you want to protected</param>
        /// <param name="stablePorcessID">stable process's ID(csrss.exe | system.exe | winlogon.exe)</param>
        /// <param name="bFileCanBeRead">the file can be shared to read or not</param>       
        /// <returns>the file's new handle value in stable process</returns>
        public static IntPtr DupFileHandle2Process(string filePath, int stablePorcessID, bool bFileCanBeRead)
        {           
            IntPtr hFile;
            SafeFileHandle fileSafehandle = null;
            IntPtr hProcess = IntPtr.Zero;
            IntPtr ipNewHandle = IntPtr.Zero; // we don't care about this, the file's new handle value in stable process                    

            // get the file handle
            hFile = NativeMethods.CreateFile(
                filePath,
                (uint)NativeMethods.ACCESS_TYPE.GENERIC_READ,
                (bFileCanBeRead?(uint)NativeMethods.FILESHARE_TYPE.FILE_SHARE_READ:0),
                IntPtr.Zero,
                (uint)NativeMethods.CREATION_TYPE.OPEN_EXISTING,
                (uint)NativeMethods.FILE_ATTRIBUTE_TYPE.FILE_ATTRIBUTE_NORMAL,
                IntPtr.Zero
              );

            fileSafehandle = new SafeFileHandle(hFile, true);
            if (!fileSafehandle.IsInvalid)
            {
                // get stable process's Handle, which will be used in DuplicateHandle function as target process parameter
                hProcess = NativeMethods.OpenProcess(NativeMethods.ProcessAccessFlags.DupHandle, false, stablePorcessID);
                if (hProcess == IntPtr.Zero)
                {
                    NativeMethods.CloseHandle(hFile);
                }
                else
                {                   
                    // Call duplicateHandle, 
                    // because we open file in this app(hFile=CreateFile..), GetCurrentProcess()'s return value, a process handle, will involve the file Handle
                    NativeMethods.DuplicateIpHandle(NativeMethods.GetCurrentProcess(), hFile, hProcess, out ipNewHandle, 0, false, 0x00000002);
                }
            }
            else
            {
                Marshal.ThrowExceptionForHR(Marshal.GetHRForLastWin32Error());
            }            

            NativeMethods.CloseHandle(hFile);
            NativeMethods.CloseHandle(hProcess);  

            return ipNewHandle;
        }
    }
}
